/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/lite/kernels/internal/reference/integer_ops/logistic.h"

#include "tensorflow/lite/c/builtin_op_data.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/quantization_util.h"
#include "tensorflow/lite/kernels/internal/reference/logistic.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/op_macros.h"
#include "tensorflow/lite/micro/kernels/kernel_util.h"
#include "cmsis/CMSIS/Core/Include/cmsis_compiler.h"
namespace tflite {
namespace ops {
namespace micro {
namespace activations {
namespace {
constexpr int kInputTensor = 0;
constexpr int kOutputTensor = 0;

#undef IS_TESTING_LUT	
	
struct OpData {
#ifdef IS_TESTING_LUT
  int32_t input_zero_point;
  int32_t input_range_radius;
  int32_t input_multiplier;
  int input_left_shift;
#endif  
  float iScale;
  float oScale;
  int iZeroPt;
  int oZeroPt;
};

TfLiteStatus CalculateArithmeticOpData(TfLiteContext* context, TfLiteNode* node,
                                       OpData* data) {
  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
  TF_LITE_ENSURE(context, input != nullptr);
  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
  TF_LITE_ENSURE(context, output != nullptr);

  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
  if (input->type == kTfLiteInt8) {
    TF_LITE_ENSURE_EQ(context, output->params.zero_point,
                      std::numeric_limits<int8_t>::min());
	#ifdef IS_TESTING_LUT
    static constexpr int kInputIntegerBits = 4;
    const double input_real_multiplier =
        static_cast<double>(input->params.scale) *
        static_cast<double>(1 << (31 - kInputIntegerBits));

    data->input_zero_point = input->params.zero_point;

    const double q = std::frexp(input_real_multiplier, &data->input_left_shift);
    data->input_multiplier = static_cast<int32_t>(TfLiteRound(q * (1ll << 31)));

    data->input_range_radius =
        CalculateInputRadius(kInputIntegerBits, data->input_left_shift, 31);
	#endif
	data->iScale = input->params.scale;
	data->oScale = output->params.scale;
	data->iZeroPt = input->params.zero_point;
	data->oZeroPt = output->params.zero_point;
  }
  return kTfLiteOk;
}
}  // namespace

void* LogisticInit(TfLiteContext* context, const char* buffer, size_t length) {
  TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
  return context->AllocatePersistentBuffer(context, sizeof(OpData));
}

TfLiteStatus LogisticPrepare(TfLiteContext* context, TfLiteNode* node) {
  TFLITE_DCHECK(node->user_data != nullptr);
  OpData* data = static_cast<OpData*>(node->user_data);

  return CalculateArithmeticOpData(context, node, data);
}

// rocky: Usually, uniform tab with grid size of 1/16 + linear interploation is precise enough
// dual tab will set grid size to 1/32 for [0, 2] and 1/8 for [0,8] (effective in [2.125,8])
#ifdef LUT_DUAL_TAB 
// 2 look up tables (LUTs)
// tabWide: x=[8], step 1, 1 elements (place holder*/

// [0, 2.0], step 1/32
static const float s_i8TabNarrow[65] = {0.5,0.507811864,0.515619916,0.523420349,0.531209373,0.538983221,0.546738152,0.554470465,0.562176501,0.569852651,0.577495365,0.585101154,0.5926666,0.600188359,0.60766317,0.615087856,0.622459331,0.629774607,0.637030794,0.644225106,0.651354865,0.658417501,0.665410559,0.672331699,0.679178699,0.685949455,0.692641983,0.699254421,0.705785028,0.712232184,0.718594393,0.724870276,0.731058579,0.737158163,0.743168009,0.749087213,0.754914987,0.760650653,0.766293643,0.771843498,0.777299861,0.782662479,0.787931196,0.793105951,0.798186778,0.803173796,0.808067214,0.812867318,0.817574476,0.822189131,0.826711794,0.831143048,0.835483537,0.839733968,0.843895103,0.847967758,0.851952802,0.855851147,0.859663751,0.86339161,0.86703576,0.870597268,0.874077235,0.877476787,0.880797078,
};
// [0.0, 8.0], step 1/8
static const float s_i8TabWide[65] = {0.5,0.531209373,0.562176501,0.5926666,0.622459331,0.651354865,0.679178699,0.705785028,0.731058579,0.754914987,0.777299861,0.798186778,0.817574476,0.835483537,0.851952802,0.86703576,0.880797078,0.893309406,0.904650535,0.914900955,0.92414182,0.932453309,0.93991335,0.94659667,0.952574127,0.957912272,0.962673113,0.966914022,0.970687769,0.974042643,0.97702263,0.979667647,0.98201379,0.984093608,0.985936373,0.987568349,0.989013057,0.990291524,0.991422515,0.992422759,0.993307149,0.994088931,0.994779874,0.995390428,0.995929862,0.996406397,0.996827317,0.997199073,0.997527377,0.997817284,0.998073265,0.998299278,0.998498818,0.998674978,0.99883049,0.998967769,0.999088949,0.999195914,0.99929033,0.999373666,0.999447221,0.999512143,0.999569443,0.999620015,0.99966465,
};
#else
// [0.0, 8.0], step 1/16
static const float s_i8TabUniform[129] = {
0.5,0.515619916,0.531209373,0.546738152,0.562176501,0.577495365,0.5926666,0.60766317,0.622459331,0.637030794,0.651354865,0.665410559,0.679178699,0.692641983,0.705785028,0.718594393,0.731058579,0.743168009,0.754914987,0.766293643,0.777299861,0.787931196,0.798186778,0.808067214,0.817574476,0.826711794,0.835483537,0.843895103,0.851952802,0.859663751,0.86703576,0.874077235,0.880797078,0.887204594,0.893309406,0.899121377,0.904650535,0.909907006,0.914900955,0.919642531,0.92414182,0.928408801,0.932453309,0.936285006,0.93991335,0.943347575,0.94659667,0.949669367,
0.952574127,0.95531913,0.957912272,0.960361161,0.962673113,0.964855154,0.966914022,0.968856169,0.970687769,0.972414718,0.974042643,0.97557691,0.97702263,0.978384667,0.979667647,0.980875963,0.98201379,0.983085087,0.984093608,0.985042913,0.985936373,0.986777178,0.987568349,0.988312742,0.989013057,0.989671847,0.990291524,0.990874363,0.991422515,0.991938008,0.992422759,0.992878571,0.993307149,0.993710098,
0.994088931,0.994445075,0.994779874,0.995094594,0.995390428,0.995668498,0.995929862,0.996175516,0.996406397,0.996623388,0.996827317,0.997018967,0.997199073,0.997368326,0.997527377,0.997676837,0.997817284,0.997949256,0.998073265,0.998189789,0.998299278,0.998402155,0.998498818,0.998589642,0.998674978,0.998755157,0.99883049,0.998901269,0.998967769,0.999030248,0.999088949,0.999144099,0.999195914,0.999244594,0.99929033,0.999333298,0.999373666,0.999411591,0.999447221,0.999480695,0.999512143,0.999541687,0.999569443,0.999595519,0.999620015,0.999643029,0.99966465,
};
#endif

#define ALWAYS_INLINE inline __attribute__((always_inline))
int ALWAYS_INLINE fast_ceilf(float x)
{
    int i;
    x += 0.999999f;
    asm volatile (
            "vcvt.S32.f32  %[r], %[x]\n"
            : [r] "=t" (i)
            : [x] "t"  (x));
    return i;
}

static int ALWAYS_INLINE fast_roundf(float x)
{
    int i;
    asm volatile (
            "vcvtr.s32.f32  %[r], %[x]\n"
            : [r] "=t" (i)
            : [x] "t"  (x));
    return i;
}

void Rky_LogisticInt8(int32_t input_zero_point, float input_scale,
                     int32_t output_zero_point, float output_scale,
                     int32_t flat_size, const int8_t* input_data,
                     int8_t* output_data)
{
    int i, ndx, iY;
    float fNdx;
    float x, y;
    float outScaleInv = 1 / output_scale;
	int8_t* output_data0 = output_data;
	uint32_t isInv;
	#ifdef IS_TESTING_LUT
		volatile int err;
		volatile int errMax = 0, errAcc = 0, valAcc = 0;
	#endif
    for (i=0; i< flat_size; i++) {
        x = ((int)(*input_data++) - input_zero_point) * input_scale;
		isInv = 0;
		if (x < 0) {
			x = -x;
			isInv = 1;
		}
		#ifdef LUT_DUAL_TAB
			if (x < 2.0f) {
				fNdx = x * 32.0f; /* div by 1/32 */
				ndx = (int)(fNdx);
				y = s_i8TabNarrow[ndx]; // sparse value
				// fine tune interpolation
				y += (s_i8TabNarrow[ndx+1] - y) * (fNdx - ndx);
			} else {
				if (x < 8.0f) {
					fNdx = x * 8.0f; /* div by 1/8 */
					ndx = (int)(fNdx);
					y = s_i8TabWide[ndx]; // sparse value
					// fine tune interpolation
					y += (s_i8TabWide[ndx+1] - y) * (fNdx - ndx);
				} else {
					y = 0.999999f;
				}
				
			}
		#else
			if (x < 8.0f) {
				fNdx = x * 16.0f; /* div by 1/16 */
				ndx = (int)(fNdx);
				y = s_i8TabUniform[ndx]; // sparse value
				// fine tune interpolation
				y += (s_i8TabUniform[ndx + 1] - y) * (fNdx - ndx);
			} else {
				y = 1.0f;
			}

		#endif
		if (isInv) {
			y = 1.0f - y;
		}
        y *= outScaleInv;
        iY = fast_roundf(y) + output_zero_point;
        iY = __SSAT(iY, 8);
		#ifdef IS_TESTING_LUT
			valAcc += *output_data > 0 ? *output_data : -*output_data;
			// *output_data = (int8_t) iY;
		    err = *output_data - (int8_t)iY;
			err = err < 0 ? -err : err;
			if (err != 0) {
				err = err;
			}
			errAcc += err;
			errMax = err > errMax ? err : errMax;
			output_data++;
		#else
			*output_data++ = (int8_t) iY;
		#endif		
    }
	#ifdef IS_TESTING_LUT
		// in case get optimized out
		err = err;
		errMax = errMax;
		errAcc = errAcc;
		valAcc = valAcc;
		err = err;
	#endif
}

TfLiteStatus LogisticEval(TfLiteContext* context, TfLiteNode* node) {
  const TfLiteEvalTensor* input =
      tflite::micro::GetEvalInput(context, node, kInputTensor);
  TfLiteEvalTensor* output =
      tflite::micro::GetEvalOutput(context, node, kOutputTensor);

  TFLITE_DCHECK(node->user_data != nullptr);
  OpData* data = static_cast<OpData*>(node->user_data);

  if (input->type == kTfLiteFloat32) {
    switch (output->type) {
      case kTfLiteFloat32: {
        reference_ops::Logistic(tflite::micro::GetTensorShape(input),
                                tflite::micro::GetTensorData<float>(input),
                                tflite::micro::GetTensorShape(output),
                                tflite::micro::GetTensorData<float>(output));
        return kTfLiteOk;
      }
      default:
        TF_LITE_KERNEL_LOG(context, "Input %s, output %s not supported.",
                           TfLiteTypeGetName(input->type),
                           TfLiteTypeGetName(output->type));
        return kTfLiteError;
    }
  } else if (input->type == kTfLiteInt8) {
    switch (output->type) {
      case kTfLiteInt8: {
		#ifdef IS_TESTING_LUT
        reference_integer_ops::Logistic(
            data->input_zero_point, data->input_range_radius,
            data->input_multiplier, data->input_left_shift,
            NumElements(input->dims),
            tflite::micro::GetTensorData<int8_t>(input),
            tflite::micro::GetTensorData<int8_t>(output));
		#endif
        Rky_LogisticInt8(
            data->iZeroPt, data->iScale,
            data->oZeroPt, data->oScale,
            NumElements(input->dims),
            tflite::micro::GetTensorData<int8_t>(input),
            tflite::micro::GetTensorData<int8_t>(output));
        return kTfLiteOk;
      }
      default:
        TF_LITE_KERNEL_LOG(context, "Input %s, output %s not supported.",
                           TfLiteTypeGetName(input->type),
                           TfLiteTypeGetName(output->type));
        return kTfLiteError;
    }
  } else {
    // TODO(b/141211002): Also support other data types once we have supported
    // temporary tensors in TFLM.
    TF_LITE_KERNEL_LOG(context, "Input %s, output %s not supported.",
                       TfLiteTypeGetName(input->type),
                       TfLiteTypeGetName(output->type));
    return kTfLiteError;
  }
  return kTfLiteOk;
}

}  // namespace activations

TfLiteRegistration Register_LOGISTIC() {
  return {/*init=*/activations::LogisticInit,
          /*free=*/nullptr,
          /*prepare=*/activations::LogisticPrepare,
          /*invoke=*/activations::LogisticEval,
          /*profiling_string=*/nullptr,
          /*builtin_code=*/0,
          /*custom_name=*/nullptr,
          /*version=*/0};
}
}  // namespace micro
}  // namespace ops
}  // namespace tflite
